fix(rocm): make CPU/Hybrid MoE graph replay safe - #378

Thanks for building this — we filed #350 and did not expect a fix, let alone one validated on the
same card we found it on. Your framing of the fix as fail-closed — disable capture and
continue on the correct eager path when the native handshake does not pass a real
capture/instantiate/replay probe — is the property we would have asked for and did not know how to
ask for. A silent wrong answer converted into a slower right one is exactly the right trade for
this class of bug.

Two things below: a ROCm datapoint your HIP memop work bears on directly, and an offer.

1. Your HIP memop port has a second consumer, and it is the half that is still private

Reading the diff, the thing that struck us is that the HIP hipStreamWriteValue64 /
hipStreamWaitValue64 / hipStreamBatchMemOp port is already here — we had it filed on our
own side as a piece of work someone would eventually have to do. It matters to us because the CPU
MoE executor is not the only thing in this engine that wants that handshake.

#311's disk-streamed PLE table resolves its stream memops by dlopening libcuda.so.1. On
ROCm that dlopen fails, the probe returns false, and the table falls back to launch-gating: a
blocking D2H of the sampled token that drains the stream, then hash, io_uring read, H2D and launch
on the host, with the GPU idle across the whole stretch. Measured on our box — half-box (one
R9700, TP=1), on tree c4b96fc1 (upstream 4b94bdc3 plus our patches 0001..0009), with
--ple-backend the only flag changed between the arms:

n is post-filter: replies are dropped if their zlib compression ratio falls below 0.40, which is
how we catch a degenerate repeated-token generation before it flatters a decode number. Note the
100k pair rests on two samples per arm — we quote it because it agrees with the 10k pair, not
as an independent result.

−27.3 % at 10k, −26.2 % at 100k, roughly 14 ms per token, against the −2.6 % … +0.2 %
band #311 reports on H100 and RTX PRO 6000 Blackwell — which was measured on the wait-sync
path. Prefill is unaffected: TTFT +0.20 s at 10k and −0.22 s at 100k, same half-box arms —
though every prefill number from this box carries "gfx1201, BLOCK_N clamped to 32" (see
Platform), and block_n = 64 does not run here at all — it raises OutOfResources before
executing an instruction — so the clamp's cost is unmeasurable rather than zero. The memory win is
real and large — on the same arms the PLE table stops pinning 47.7 GiB of host RAM — which is
why we run it anyway.

We can only bound the split, not measure it. Sampled from /proc/diskstats on the same
half-box disk arm, during decode the NVMe served 300 reads/s at 0.648 ms average = 16 reads per
token, one 4K O_DIRECT read per row. Fully serialised that is 10.4 ms; if the store's io_uring
batch puts all 16 in flight together it is ~0.65 ms. /proc/diskstats cannot tell those apart, so
the disk accounts for 0.65–10.4 ms and the launch-gating stall for the remaining 4–13.7
ms. What the bound does establish is that the disk alone cannot reach 14 ms even in its worst
case. Splitting it properly needs a wall-clock inside the fill, which we have not written — so
treat the split as bounded, not measured.

Where your PR lands on it: the eager half is now reachable — on ROCm your module-level
memops_probe / memop_submit / memop_sync resolve to the HIP implementations, so a second
consumer could stop dlopening libcuda and just call them. The half that is still
CpuMoeExecutor-owned is the capture half — and that is the half #311 needs, since by its own
description the decode graph launches first and waits on the flag right before consuming the rows.
If the ROCm graph-node splicing ever becomes shareable, there is a caller waiting for it.

And the ROCm-7.14 fact you documented in that file — that ordinary hipStreamWrite/WaitValue
and hipStreamBatchMemOp calls execute during capture but are not recorded in the graph, so
you add explicit batch-memory-op nodes and splice them into the dependency set — is worth flagging
loudly to anyone attempting a second port. Someone who ports the eager calls, sees them work, and
captures a graph gets silently stale values on replay: the same failure shape as #350 itself, in a
different file.

None of this is a request. It is a note that the two land in the same place.

2. An offer against your draft-exit criteria

You list "reproduce the replay suite on RDNA3", "complete repeated graph rebuild/replay and
long-running serving stability checks", and NVIDIA CI for the preserved CUDA path.

We can help with parts of the middle one, on RDNA4 rather than RDNA3. We run two Radeon AI
PRO R9700 (gfx1201) on Linux under ROCm 7.14 / PyTorch 2.11, serving a 100 GiB-class MoE
continuously, and we can run your targeted safety and replay suite
(tests/engine/test_cpu_moe_graph_safety.py and the tests/moe/ graph-replay cases) on both
cards and report the numbers the way we did on #132.

We can also exercise it at TP=2. To be exact about what is new there: we have reported TP=2
serving before, on #135, so the gap is not TP=2 as such — it is that nobody has run the
CPU/Hybrid graph-replay path at TP>1, on any card. That is worth closing, because from our own
experience a caught per-rank exception behaves very differently at TP=2 than at TP=1. That
difference is a scheduler problem rather than yours, and we have written it up separately.

Three honest caveats before you count on any of it.

We cannot give you the "long-running serving stability" half. We serve --moe-backend offload, not cpu or hybrid, so a serving soak on this code path is not something our
production traffic would ever produce — it would be a rig we stand up specially, and it would not
be representative of a real workload the way our offload row is. Repeated rebuild/replay and the
suite are the parts we can do well.

We cannot promise a date. The box is a shared machine with a live soak on it. Say what you
want run and we will fit it in.

It is a second gfx1201 datapoint, not a clean-main CI result. RDNA4, on a tree that is
upstream 4b94bdc3 plus a local patch series, with the BLOCK_N clamp in the Platform note
below.

No quality claim from us in any case — we have no fidelity instrument on this box, which is why
#350 was reported as a numerical-deviation finding with a relative error rather than as a quality
regression.

Platform

2 × AMD Radeon AI PRO R9700 (gfx1201, 32 GiB each), Linux, ROCm 7.14, PyTorch 2.11, 6 physical
cores / 122 GiB host RAM. Serving qwen4_exp / Qwen3.8-Flash-Next NVFP4 with --moe-backend offload; the CPU/Hybrid finding in #350 came from running this repo's kernel suite on gfx1201,
not from this serving configuration.

Our tree is upstream 4b94bdc3 plus a local patch series, whose build asserts tree c4b96fc1
for the 0001..0009 stack every §1 number was measured on. The QSA prefill tile is halved until 2 · head_dim · block_n · itemsize fits a 60 KiB LDS budget (#349) — at this checkpoint's head_dim = 256 in bf16 that lands on BLOCK_N = 32. So every prefill number here is on the clamped tile
and is not comparable to a CUDA run of the same engine.

Box labelling, since this box has two cards: every number in §1 — decode, TTFT, the 47.7 GiB
and the /proc/diskstats figures alike — is half-box, one R9700 at TP=1, on tree c4b96fc1
(upstream 4b94bdc3 plus our patches 0001..0009), with --ple-backend the only flag changed
between the arms. Anything we quote at --tp-size 2 is full-box. Every number is named
against the tree it was measured on, and none of them is measured on upstream main.